import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


fidx_u_16 = [0,0,6,0,0,1,1,4,5,1,3,0,0,0,2,3] # 论文所选频率的水平指数
fidx_v_16 = [0,1,0,5,2,0,2,0,0,6,0,4,6,3,2,5] # 论文所选频率的垂直指数
# 获取一维通道的DCT
# L为输入的长度
def get_1d_dct(pos, freq, L):
    result = math.cos(math.pi * freq * (pos + 0.5) / L) / math.sqrt(L)
    if freq == 0:
        return result
    else:
        return result * math.sqrt(2)

# 推广到二维DCT
def get_dct_weights(input_width, input_height, input_channel, fidx_u= fidx_u_16, fidx_v= fidx_v_16):
    # 论文将频率空间分为7*7,实验结果表明,[u,v]=[0,0],即GAP的效果最好(乐),当然了，改进的点在于在信道注意力中使用多个频率分量
    # 使不同大小的频率与 7x7 频率空间相同
    fidx_horizon = [u*(input_width//7) for u in fidx_u]
    fidx_vertical = [v*(input_height//7) for v in fidx_v]
    dct_weights = torch.zeros(1, input_channel, input_width, input_height)
    c_part = input_channel // len(fidx_horizon)
    # 拆分通道
    for i, (u_x, v_y) in enumerate(zip(fidx_horizon, fidx_vertical)):
        for t_x in range(input_width):
            for t_y in range(input_height):
                dct_weights[:, i * c_part: (i+1)*c_part, t_x, t_y]\
                =get_1d_dct(t_x, u_x, input_width) * get_1d_dct(t_y, v_y, input_height)

    return dct_weights

# SE-Block单元--SEblock是一个子结构，几乎可以嵌入任何一个神经网络模型之中
class FcaBlock(nn.Module):
    def __init__(self, input_channel, width, height, reduction=16):
        super(FcaBlock, self).__init__()
        self.width = width
        self.height = height


        # self.adaptive_avg_pool = nn.AdaptiveAvgPool2d(1)  # 全局自适应池化
        self.register_buffer('pre_computed_dct_weights', get_dct_weights(self.width, self.height, input_channel))
        self.fc = nn.Sequential(
            nn.Linear(input_channel, input_channel // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(input_channel // reduction, input_channel),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, h, w = x.size()

        # squeeze操作:(b,c,h,w)->(b,c)
        #y = self.adaptive_avg_pool(x).view(b, c)
        y = F.adaptive_avg_pool2d(x, (self.height, self.width))
        y = torch.sum(y*self.pre_computed_dct_weights, dim=[2,3])

        # FC获取通道注意力权重，是具有全局信息的
        y = self.fc(y).view(b, c, 1, 1)

        # 注意力作用每一个通道上
        y = x * y.expand_as(x)
        # 残差连接
        return x+y


class _DenseLayer(nn.Module):
    def __init__(self, num_input_features, growth_rate, bn_size, drop_rate=0,width=32,height=32,reduction=16):
        super(_DenseLayer, self).__init__()
        self.drop_rate = drop_rate
        self.dense_layer = nn.Sequential(
            nn.BatchNorm2d(num_input_features),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=num_input_features, out_channels=bn_size * growth_rate, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(bn_size * growth_rate),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=bn_size * growth_rate, out_channels=growth_rate, kernel_size=3, stride=1, padding=1, bias=False)
        )
        self.fca = self.fca = FcaBlock(growth_rate, width, height, reduction)
        self.dropout = nn.Dropout(p=self.drop_rate)

    def forward(self, x):
        y = self.dense_layer(x)
        y = self.fca(y)
        if self.drop_rate > 0:
            y = self.dropout(y)

        return torch.cat([x, y], dim=1)


class _DenseBlock(nn.Module):
    def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate=0):
        super(_DenseBlock, self).__init__()
        layers = []
        for i in range(num_layers):
            layers.append(_DenseLayer(num_input_features + i * growth_rate, growth_rate, bn_size, drop_rate))
        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        return self.layers(x)


class _TransitionLayer(nn.Module):
    def __init__(self, num_input_features, num_output_features):
        super(_TransitionLayer, self).__init__()
        self.transition_layer = nn.Sequential(
            nn.BatchNorm2d(num_input_features),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=num_input_features, out_channels=num_output_features, kernel_size=1, stride=1, padding=0, bias=False),
            nn.AvgPool2d(kernel_size=2, stride=2)
        )

    def forward(self, x):
        return self.transition_layer(x)


class DenseNet(nn.Module):
    def __init__(self, num_init_features=64, growth_rate=32, blocks=(6, 12, 24, 16), bn_size=4, drop_rate=0, num_classes=1000):
        super(DenseNet, self).__init__()

        self.features = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=num_init_features, kernel_size=7, stride=2, padding=3, bias=False),
            nn.BatchNorm2d(num_init_features),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )

        num_features = num_init_features
        self.layer1 = _DenseBlock(num_layers=blocks[0], num_input_features=num_features, growth_rate=growth_rate, bn_size=bn_size, drop_rate=drop_rate)
        num_features = num_features + blocks[0] * growth_rate
        self.transtion1 = _TransitionLayer(num_input_features=num_features, num_output_features=num_features // 2)

        num_features = num_features // 2
        self.layer2 = _DenseBlock(num_layers=blocks[1], num_input_features=num_features, growth_rate=growth_rate, bn_size=bn_size, drop_rate=drop_rate)
        num_features = num_features + blocks[1] * growth_rate
        self.transtion2 = _TransitionLayer(num_input_features=num_features, num_output_features=num_features // 2)

        num_features = num_features // 2
        self.layer3 = _DenseBlock(num_layers=blocks[2], num_input_features=num_features, growth_rate=growth_rate, bn_size=bn_size, drop_rate=drop_rate)
        num_features = num_features + blocks[2] * growth_rate
        self.transtion3 = _TransitionLayer(num_input_features=num_features, num_output_features=num_features // 2)

        num_features = num_features // 2
        self.layer4 = _DenseBlock(num_layers=blocks[3], num_input_features=num_features, growth_rate=growth_rate, bn_size=bn_size, drop_rate=drop_rate)
        num_features = num_features + blocks[3] * growth_rate

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(num_features, num_classes)

    def forward(self, x):
        x = self.features(x)

        x = self.layer1(x)
        x = self.transtion1(x)
        x = self.layer2(x)
        x = self.transtion2(x)
        x = self.layer3(x)
        x = self.transtion3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, start_dim=1)
        x = self.fc(x)

        return x

# model =DenseNet(blocks=(6, 12, 32, 32), num_classes=100).to(device)
#
# summary(model, input_size=(3, 32, 32))